# train.py
import pandas as pd
import torch
from tqdm import tqdm

from ..config.config import Config
from ..model.CasreModel import CasRel

import json

from ..until.data_process import TransferData

class CasrelModelTrain:
    def __init__(self):
        self.config = Config()
        self.casrel_model = CasRel()
        self.transfer_data = TransferData()
    def model2train(self, train_iter, dev_iter, optimizer, ):
        best_triple_f1 = 0
        for epoch in range(self.config.epochs):
            self.train_epoch( train_iter, dev_iter, optimizer, best_triple_f1, epoch)
        torch.save(self.casrel_model.state_dict(), f'{self.config.save_model_path}\\last_model.pth')


    def train_epoch(self,model, train_iter, dev_iter, optimizer, best_triple_f1, epoch):
        for step, (inputs, labels) in enumerate(tqdm(train_iter)):
            self.casrel_model.train()
            logist = self.casrel_model(**inputs)
            loss = self.casrel_model.compute_loss(**logist, **labels)

            model.zero_grad()
            loss.backward()
            optimizer.step()

            if step % 1500 == 0:
                torch.save(self.casrel_model.state_dict(),
                           f'{self.config.save_model_path}\\epoch_%s_model_%s.pth' % (epoch, step))
                results = self.model2dev(dev_iter)
                print(results[-1])
                if results[-2] > best_triple_f1:
                    best_triple_f1 = results[-2]
                    torch.save(self.casrel_model.state_dict(), f'{self.config.save_model_path}\\best_f1.pth')
                    print('epoch:{},'
                          'step:{},'
                          'sub_precision:{:.4f}, '
                          'sub_recall:{:.4f}, '
                          'sub_f1:{:.4f}, '
                          'triple_precision:{:.4f}, '
                          'triple_recall:{:.4f}, '
                          'triple_f1:{:.4f},'
                          'train loss:{:.4f}'.format(epoch,
                                                     step,
                                                     results[0],
                                                     results[1],
                                                     results[2],
                                                     results[3],
                                                     results[4],
                                                     results[5],
                                                     loss.item()))

        return best_triple_f1


    def model2dev(self, dev_iter):
        '''
        验证模型效果
        :param model:
        :param dev_iter:
        :return:
        '''
        self.casrel_model.eval()
        # 定义一个df，来展示模型的指标。
        df = pd.DataFrame(columns=['TP', 'PRED', "REAL", 'p', 'r', 'f1'], index=['sub', 'triple'])
        df.fillna(0, inplace=True)
        for inputs, labels in tqdm(dev_iter):
            logist = self.casrel_model(**inputs)
            pred_sub_heads = self.transfer_data.convert_score_to_zero_one(logist['pred_sub_heads'])
            pred_sub_tails = self.transfer_data.convert_score_to_zero_one(logist['pred_sub_tails'])
            sub_heads = self.transfer_data.convert_score_to_zero_one(labels['sub_heads'])
            sub_tails = self.transfer_data.convert_score_to_zero_one(labels['sub_tails'])
            batch_size = inputs['input_ids'].shape[0]
            obj_heads = self.transfer_data.convert_score_to_zero_one(labels['obj_heads'])
            obj_tails = self.transfer_data.convert_score_to_zero_one(labels['obj_tails'])
            pred_obj_heads = self.transfer_data.convert_score_to_zero_one(logist['pred_obj_heads'])
            pred_obj_tails = self.transfer_data.convert_score_to_zero_one(logist['pred_obj_tails'])

            for batch_index in range(batch_size):

                pred_subs = self.transfer_data.extract_sub(pred_sub_heads[batch_index].squeeze(),
                                        pred_sub_tails[batch_index].squeeze())

                true_subs = self.transfer_data.extract_sub(sub_heads[batch_index].squeeze(),
                                        sub_tails[batch_index].squeeze())

                pred_ojbs = self.transfer_data.extract_obj_and_rel(pred_obj_heads[batch_index],
                                                pred_obj_tails[batch_index])

                true_objs = self.transfer_data.extract_obj_and_rel(obj_heads[batch_index],
                                                obj_tails[batch_index])

                df['PRED']['sub'] += len(pred_subs)
                df['REAL']['sub'] += len(true_subs)

                for true_sub in true_subs:
                    if true_sub in pred_subs:
                        df['TP']['sub'] += 1

                df['PRED']['triple'] += len(pred_ojbs)
                df['REAL']['triple'] += len(true_objs)
                for true_obj in true_objs:
                    if true_obj in pred_ojbs:
                        df['TP']['triple'] += 1

        df.loc['sub', 'p'] = df['TP']['sub'] / (df['PRED']['sub'] + 1e-9)
        df.loc['sub', 'r'] = df['TP']['sub'] / (df['REAL']['sub'] + 1e-9)
        df.loc['sub', 'f1'] = 2 * df['p']['sub'] * df['r']['sub'] / (df['p']['sub'] +
                                                                     df['r']['sub'] +
                                                                     1e-9)
        sub_precision = df['TP']['sub'] / (df['PRED']['sub'] + 1e-9)
        sub_recall = df['TP']['sub'] / (df['REAL']['sub'] + 1e-9)
        sub_f1 = 2 * sub_precision * sub_recall / (sub_precision + sub_recall + 1e-9)

        df.loc['triple', 'p'] = df['TP']['triple'] / (df['PRED']['triple'] + 1e-9)
        df.loc['triple', 'r'] = df['TP']['triple'] / (df['REAL']['triple'] + 1e-9)
        df.loc['triple', 'f1'] = 2 * df['p']['triple'] * df['r']['triple'] / (
                df['p']['triple'] + df['r']['triple'] + 1e-9)

        triple_precision = df['TP']['triple'] / (df['PRED']['triple'] + 1e-9)
        triple_recall = df['TP']['triple'] / (df['REAL']['triple'] + 1e-9)
        triple_f1 = 2 * triple_precision * triple_recall / (
                triple_precision + triple_recall + 1e-9)

        return sub_precision, sub_recall, sub_f1, triple_precision, triple_recall, triple_f1, df

    # coding:utf-8
    """
    电商知识图谱项目中的关系抽取预测脚本 (predict_ecom.py)
    使用 CasRel 模型从商品文本中抽取三元组 (subject, predicate, object)
    """
    def load_model(self,model_path):
        """
        加载训练好的 CasRel 模型
        :param model_path: 模型权重文件路径
        :return: CasRel 模型实例
        """
        model = self.casrel_model.to(self.config.device)
        model.load_state_dict(torch.load(model_path, map_location=self.config.device))
        model.eval()
        return model

    def get_inputs(self,text):
        """
        将输入文本 tokenize 并生成模型所需张量
        返回 inputs 字典与模型实例
        """
        encoding = self.config.tokenizer(text, return_tensors='pt')
        input_ids = encoding['input_ids'].to(self.config.device)
        attention_mask = encoding['attention_mask'].to(self.config.device)

        seq_len = input_ids.size(1)
        sub_head2tail = torch.zeros((1, seq_len), device=self.config.device)
        sub_len = torch.ones((1,), dtype=torch.float, device=self.config.device)

        # 先预测主实体位置
        with torch.no_grad():
            encoded = self.casrel_model.get_encoded_text(input_ids, attention_mask)
            sub_heads, sub_tails = self.casrel_model.get_subs(encoded)
            heads_bin = self.transfer_data.convert_score_to_zero_one(sub_heads.squeeze(0))
            tails_bin = self.transfer_data.convert_score_to_zero_one(sub_tails.squeeze(0))
            subs = self.transfer_data.extract_sub(heads_bin, tails_bin)
            if subs:
                h, t = subs[0]
                sub_head2tail[0, h:t + 1] = 1
                sub_len[0] = t - h + 1

        inputs = {
            'input_ids': input_ids,
            'mask': attention_mask,
            'sub_head2tail': sub_head2tail,
            'sub_len': sub_len
        }
        return inputs

    def predict_triples(self,text, rel_map_path):
        """
        对单条文本进行三元组预测
        :param text: 待预测文本
        :param model: CasRel 模型实例
        :param rel_map_path: JSON 文件，关系 id->名称 映射
        :return: dict, 包含原文和抽取的 spo_list
        """
        # 加载关系映射
        with open(rel_map_path, 'r', encoding='utf-8') as f:
            rel_id2name = json.load(f)

        inputs = self.get_inputs(text)
        with torch.no_grad():
            outputs = self.casrel_model(**inputs)

        sub_heads = self.transfer_data.convert_score_to_zero_one(outputs['pred_sub_heads'].squeeze(0))
        sub_tails = self.transfer_data.convert_score_to_zero_one(outputs['pred_sub_tails'].squeeze(0))
        obj_heads = self.transfer_data.convert_score_to_zero_one(outputs['pred_obj_heads'].squeeze(0))
        obj_tails = self.transfer_data.convert_score_to_zero_one(outputs['pred_obj_tails'].squeeze(0))

        # Token 转回文本
        tokens = self.config.tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
        sentence = ''.join(tokens[1:-1])

        subs = self.transfer_data.extract_sub(sub_heads, sub_tails)
        objs = self.transfer_data.extract_obj_and_rel(obj_heads[0], obj_tails[0])

        spo_list = []
        if subs and objs:
            # 如果 obj 数量更多，则重复主实体
            if len(objs) > len(subs):
                subs = subs * len(objs)
            for (h, t), (rel_id, oh, ot) in zip(subs, objs):
                subj = ''.join(tokens[h: t + 1])
                obj = ''.join(tokens[oh: ot + 1])
                predicate = rel_id2name.get(str(rel_id), 'unknown')
                spo_list.append({
                    'subject': subj,
                    'predicate': predicate,
                    'object': obj
                })

        return {'text': sentence, 'spo_list': spo_list}
